#include <stdio.h>

#include <cfloat>  // 标准C++浮点数限制头文件

#include "common/config.h"
#include "gpu/kernel.h"

__global__ void calc_esdf_row(int map_size, int map_col_size, float* tmp_buffer, int* obs_index_data,
                              int* obs_start_index_data) {  // 行扫描
    int idx = threadIdx.x + blockIdx.x * blockDim.x;        // 处理第idx个像素
    if (idx >= map_size) {
        return;
    }

    int row = idx / map_col_size;
    int col = idx % map_col_size;

    int obs_start_index = obs_start_index_data[row];
    int obs_end_index = obs_start_index_data[row + 1];

    float min_dist = FLT_MAX;
    for (int i = obs_start_index; i < obs_end_index; i++) {
        int obs_col = obs_index_data[i];
        float dist = (obs_col - col) * (obs_col - col);
        min_dist = min(min_dist, dist);
    }

    tmp_buffer[idx] = min_dist;

    // if (row == 10 && (col <= 450 && col >= 448)) {
    //     printf("row: %d, col: %d, min_dist: %.3f, obs_start_index: %d, obs_end_index: %d\n", row, col, min_dist, obs_start_index, obs_end_index);
    // }

    // if (row == 11 && (col >= 440 && col <= 455)) {
    //     printf("row: %d, col: %d, min_dist: %.3f, obs_start_index: %d, obs_end_index: %d\n", row, col, min_dist, obs_start_index, obs_end_index);
    // }
    return;
}

__global__ void calc_esdf_col(int map_size,
                              int map_row_size,
                              int map_col_size,
                              float* tmp_buffer,         // 行扫描的结果
                              float* distance_buffer) {  // 列扫描
    int idx = threadIdx.x + blockIdx.x * blockDim.x;     // 处理第idx个像素
    if (idx >= map_size) {
        return;
    }

    int row = idx / map_col_size;
    int col = idx % map_col_size;

    int row_start = max(row - kRowIndexRange, 0);
    int row_end = min(row + kRowIndexRange, map_row_size - 1);
    float min_dist = FLT_MAX;
    for (int i = row_start; i <= row_end; i++) {
        float temp_distance = tmp_buffer[i * map_col_size + col];
        float val = (i - row) * (i - row) + temp_distance;
        min_dist = min(min_dist, val);
    }
    distance_buffer[idx] = sqrt(min_dist) * kResolution;

    // if (col == 25 && row == 50) {
    //     printf("row: %d, col: %d, row_start: %d, row_end: %d, min_dist: %.2f\n", row, col, row_start, row_end, min_dist);
    //     for (int i = 0; i <= row_end; i++) {
    //         float temp_distance = tmp_buffer[i * map_col_size + col];
    //         printf("(%d, %.2f)", i, sqrt(temp_distance));
    //     }
    //     printf("\n");
    // }

    return;
}

void calc_esdf_map_launcher(cudaStream_t* stream,
                            int map_row_size,
                            int map_col_size,
                            float resolution,
                            uint8_t* occupancy_buffer_inflate,
                            float* tmp_buffer,
                            float* distance_buffer,
                            int* obs_index_data,
                            int* obs_start_index_data) {
    dim3 block_size(256);
    int map_size = map_col_size * map_row_size;
    dim3 grid_size(static_cast<int>((map_size + block_size.x) / block_size.x));

    calc_esdf_row<<<grid_size, block_size, 0, *stream>>>(map_size, map_col_size, tmp_buffer, obs_index_data, obs_start_index_data);
    calc_esdf_col<<<grid_size, block_size, 0, *stream>>>(map_size, map_row_size, map_col_size, tmp_buffer, distance_buffer);
}